import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

_ = load_dotenv(find_dotenv())

api_key = os.environ.get('ZHIPU_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")

model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

prompt_template = ChatPromptTemplate.from_template("请写一首关于{topic}的诗，并需要有标题。")
chat_model = ChatZhipuAI(api_key=api_key,
                  model=model,
                  temperature=0.9,
                  stream=True,            
                )
output_parser = StrOutputParser()

chain = prompt_template | chat_model | output_parser

for chunk in chain.stream("春天"):
    print(chunk, end="", flush=True)
